import numpy as np
import matplotlib.pyplot as plt

def particle_swarm_algorithm(f, r, n, esp):
    xs = np.random.rand(n)*(r[1]-r[0]) + r[0]
    vs = np.ones(n)

    xs_best = xs.copy()
    ys_best = f(xs_best)

    def get_best(ys, xs):
        i = ys.argsort()[-1]
        return (ys[i], xs[i])

    (x_global, y_global) = get_best(ys_best, xs_best)
    
    (w, rr, c1, c2) = (0.7, 0.7, 2, 2)
    while True:
        change = 0
        for i in range(n):
            print(i)
            if np.abs(xs[i] - x_global) < esp: continue

            (xi,eta) = np.random.rand(2)

            vs[i] = w*vs[i] + c1*xi*(xs_best[i]-xs[i]) + c2*eta*(x_global-xs[i])         
            xs[i] = xs[i] + vs[i]*rr
            if xs[i] < r[0]: xs[i] = r[0]
            if xs[i] > r[1]: xs[i] = r[1]

            y = f(xs[i])
            
            if y > ys_best[i]: (xs_best[i], ys_best[i]) = (xs[i], y)
            if y > y_global: (x_global, y_global) = (xs[i], y)
            change += 1
        print(xs)
        print('change' + str(change))
        if change == 0: break
            

    return get_best(ys_best, xs_best)


def show_function(f, r, esp):
    
    x = np.linspace(r[0], r[1],1000)
    y = f(x)
    plt.plot(x, y)

    (yb,xb) = particle_swarm_algorithm(f, r, 100, esp)

    plt.plot(xb, yb, 'o')
    plt.show()



if __name__ == '__main__':
    
    f = lambda x: 1-np.cos(3*x)*np.exp(-x)
    show_function(f, [0,4], 0.000001)
